# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: jhkim

import os
from absl import app, flags

import logging
import functools
import numpy
import pandas

from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import coin.strategy.mm.tool.archive_base as abase
import coin.pnl.sim_stat_plot as ssplot


class FeedReader(object):
  def __init__(self, norm_symbol, is_index):
    self.norm_symbol = norm_symbol
    self.product = None
    self.is_index = is_index
    if is_index:
      self.columns = ['ts','indexp']
    else:
      self.columns = ['ts','askp','askq','bidp','bidq']
    self.rows = []

  def on_book_reset(self, book_builder_name, book_builder):
    book_builder.subscribe(
        self.product, functools.partial(self.on_book, self.product))
    if self.is_index:
      book_builder.subscribe_trade(
          self.product, functools.partial(self.on_index, self.product))
    else:
      book_builder.subscribe_trade(
          self.product, functools.partial(self.on_trade, self.product))

  def on_book(self, product, book):
    assert product == self.product
    self.rows.append((
        book.timestamp,
        book.get_ask_array(1)[0,0],
        book.get_ask_array(1)[0,1],
        book.get_bid_array(1)[0,0],
        book.get_bid_array(1)[0,1],
    ))

  def get_df(self, timeindex=False):
    df = pandas.DataFrame(self.rows, columns=self.columns)
    if timeindex:
      df.index = pandas.DatetimeIndex(df['ts'])
    mkt, exch, apiv, recipe, pname = self.norm_symbol.split(":")
    mea = f"{mkt}.{exch}.{apiv}"
    df['mea'] = mea
    df['product_name'] = pname
    return df

  def on_index(self, product, index):
    self.rows.append((
        index.timestamp,
        index.price
    ))

  def on_trade(self, product, trade):
    pass


def read_price(
    trading_date,
    time_range,
    norm_symbol,
    feed_machine,
    is_index):
  fr = FeedReader(norm_symbol, is_index=is_index)
  tds = abase.get_trading_dates(trading_date)
  if len(tds) > 1:
    dfs = []
    for td in tds:
      dfs.append(read_price(td.strftime("%Y%m%d"), time_range, norm_symbol, feed_machine, is_index))
    return pandas.concat(dfs, axis=0, sort=False)
  else:
    for product in abase.run_from_archive_norm(
        trading_date=trading_date,
        time_range=time_range,
        norm_symbol=norm_symbol,
        feed_machine=feed_machine,
        on_book_reset=fr.on_book_reset):
      fr.product = product
    return fr.get_df(timeindex=True)


def main(argv):
  dfidx = read_price(
      '20210501-20210511',
      '0-24',
      'Futures:Binance:v1:non_tbs:BTC-USDT.PERPETUAL',
      'feed-05.ap-northeast-1.aws',
      is_index=True)
  assert len(dfidx) > 0
  dfprc = read_price(
      '20210501-20210511',
      '0-24',
      'Futures:Binance:v1:l1_realtime_move2bp:BTC-USDT.PERPETUAL',
      'feed-05.ap-northeast-1.aws',
      is_index=False)
  dfprc['midp'] = 0.5 * (dfprc['askp'] + dfprc['bidp'])
  os.makedirs('pic', exist_ok=True)
  ssplot.setup_plt()
  plt.rcParams["figure.figsize"] = 14, 12
  plt.suptitle('Futures:Binance:v1:l1_realtime_move2bp:BTC-USDT.PERPETUAL')
  plt.subplot(211)
  plt.plot(dfprc['midp'], 'r-', lw=0.5, drawstyle='steps-post')
  plt.plot(dfidx['indexp'], 'g-', lw=0.5, drawstyle='steps-post')
  plt.legend(['midp', 'indexp'], loc='upper left')
  plt.subplot(212)
  dfconcat = pandas.concat([dfidx, dfprc], axis=0, sort=False).sort_index()
  dfconcat = dfconcat.ffill()
  dfconcat['midp_minus_indexp'] = dfconcat['midp'] - dfconcat['indexp']
  plt.plot(dfconcat['midp_minus_indexp'], 'k-', lw=0.5, drawstyle='steps-post')
  plt.legend(['midp - indexp'], loc='upper left')
  plt.savefig('pic/basis.png')
  plt.close()
  import pdb; pdb.set_trace()


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
